import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
import os
import warnings
warnings.filterwarnings('ignore')


class Dataset:
    def __init__(self, root_path='./dataset', data_path='RCLCd.csv', scale=True):
        self.root_path = root_path
        self.data_path = data_path
        self.scale = scale
        self.raw_data = pd.read_csv(os.path.join(root_path, data_path), encoding='UTF-8')
        self.__read_data__()

    def __read_data__(self):
        data = self.raw_data[['price']].values  # get data astype of ndarray
        self.inv_data = data
        stamp = self.raw_data[['date']]     # get time stamp
        # stamp = stamp.date.apply(lambda row: pd.to_datetime(row), 1)
        data_stamp = stamp.values.flatten().astype(str)

        self.year = pd.to_datetime(data_stamp).year
        self.month = pd.to_datetime(data_stamp).month
        self.day = pd.to_datetime(data_stamp).day


        # train:val:test = 10:3:3
        n = len(data) // 16
        train_size = int(10 * n)
        val_size = int(3 * n)

        train_data, self.train_stamp = data[:train_size], data_stamp[:train_size]

        if self.scale:
            self.scaler = StandardScaler()
            self.scaler.fit(train_data)
            self.train_data = self.scaler.transform(train_data)
            data = self.scaler.transform(data)

        self.data = data
        self.data_stamp = data_stamp
        self.val_data, self.val_stamp = data[train_size: train_size+val_size], data_stamp[train_size: train_size+val_size]
        self.test_data, self.test_stamp = data[train_size+val_size:], data_stamp[train_size+val_size:]

    @property
    def train(self):
        return self.train_data, self.train_stamp

    @property
    def val(self):
        return self.val_data, self.val_stamp

    @property
    def test(self):
        return self.test_data, self.test_stamp

    def transform(self, x: np.ndarray):
        """
        标准化
        """
        return self.scaler.transform(x)

    def inverse_transform(self, x: np.ndarray):
        """
        反标准化
        """
        return self.scaler.inverse_transform(x)

    def __len__(self):
        return self.data.shape[0]


class Dataset_Pred:
    def __init__(self, root_path='./dataset', data_path='all.csv', scale=True, pred_step=10):
        self.root_path = root_path
        self.data_path = data_path
        self.scale = scale
        self.raw_data = pd.read_csv(os.path.join(root_path, data_path), encoding='UTF-8')
        self.step = pred_step
        self.__read_data__()

    def __read_data__(self):
        stamp = self.raw_data[['date']]     # get time stamp
        data_stamp = stamp.values.flatten().astype(str)
        self.year = pd.to_datetime(data_stamp).year

        columns = self.raw_data.columns.drop('date')

        data = self.raw_data[columns].values
        train_data = data[:-self.step]
        if self.scale:
            self.scaler = StandardScaler()
            self.scaler.fit(train_data)
            data = self.scaler.transform(data)
        self.data = data
        self.train_data = data[: -self.step, : -1]
        self.target = data[:-self.step, -1:]
        self.pred_data = data[-self.step:, : -1]

    def transform(self, x: np.ndarray):
        """
        标准化
        """
        return self.scaler.transform(x)

    def inverse_transform(self, x: np.ndarray):
        """
        反标准化
        """
        return self.scaler.inverse_transform(x)

    def __len__(self):
        return len(self.train_data)

